from random import random, choice, randint
from collections import defaultdict
import json
from config import *

class Task():
    def __init__(self):
        self.id = None

        self.timecost = randint(1, avg_timecost * 2 - 1)
        self.data_size = randint(0, max_data_size)

        self.executor = None
        self.depth = 0

        self.time2son_task = {}   #多久之后执行了子任务

    def toLine(self):
        return f'{self.id};{self.data_size};{self.timecost};{self.depth};{json.dumps(dict(self.time2son_task))}'

    @staticmethod
    def toMeta():
        return 'id;data_size;timecost;depth;call_tree'

    def __str__(self):
        return f'{self.id}'

# 没有考虑回来的数据
def generateBenchmark(file_name, max_time = 1000, task_per_slide = 20):
    all_tasks = []

    for time in range(max_time):
        task_number = randint(0, task_per_slide)

        un_finished_task = [task for task in all_tasks if task.start_time + task.timecost > time]
        
        for _ in range(task_number):
            new_task = Task()
            new_task.start_time = time
            new_task.id = len(all_tasks)
            all_tasks.append(new_task)

            if len(un_finished_task) > 0:
                parent_task = choice(un_finished_task)
                un_finished_task.remove(parent_task)
                parent_task.time2son_task[time - parent_task.start_time] = new_task.id
                new_task.depth =  parent_task.depth + 1

    print('生成了', len(all_tasks), '个tasks')


    file = open(file_name, 'w', encoding='utf-8')
    file.write(Task.toMeta() + '\n')
    for task in all_tasks:
        file.write(task.toLine() + "\n")
    file.close()

def loadBenchmark(file_name):
    file = open(file_name, 'r', encoding='utf-8')
    lines = file.read().strip().split('\n')
    file.close()

    root_tasks = []
    id2task = {}
    for line in lines[1:]:
        elm = line.split(';')
        id, data_size, timecost, depth, call_tree = elm
        new_task = Task()

        new_task.id = int(elm[0])
        new_task.data_size = int(data_size)
        new_task.timecost = int(timecost)
        new_task.depth = int(depth)
        new_task.time2son_task = {
            int(time): task_id
            for time, task_id in json.loads(call_tree).items()
        }
        
        id2task[new_task.id] = new_task
        new_task.id2task = id2task
        if new_task.depth == 0:
            root_tasks.append(new_task)

    assert len(root_tasks) > 0 
    return root_tasks

if __name__ == "__main__":
    benchmark_name = 'bechmark100_5.txt'
    generateBenchmark(benchmark_name, max_time = 100, task_per_slide = int(node_number/avg_timecost))
    root_tasks = loadBenchmark(benchmark_name)

